# 1.导包
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier     # CART分类决策树
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report
import matplotlib.pyplot as plt
from sklearn.tree import plot_tree      # 用于绘制决策树图

# 2.读取数据
df = pd.read_csv("train.csv")
# print(df.head())

# 3.数据基本处理
# 3.1 填充年龄这一列的空值，使用均值去填充
new_df = df.fillna(df['Age'].mean())

# 3.2 删除对目标值没有价值的特征列
new_df = new_df.drop(["PassengerId", "Name", "Ticket", "Cabin", "Embarked"], axis=1)

# 3.3 将非数字类型的使用热编码进行转换
one_hot_df = pd.get_dummies(new_df)
print(one_hot_df)

# 3.4 得到特征值、目标值
x = one_hot_df.iloc[:, 1:]
y = one_hot_df.iloc[:, 0]

# 4.划分数据集
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=922, stratify=y)

# 5.特征工程
# 5.1 特征预处理
transformer = StandardScaler()
x_train = transformer.fit_transform(x_train)
x_test = transformer.fit_transform(x_test)

# 6.模型构建
"""
    criterion：设置决策树的类型
    max_depth：限定决策树的最大深度
    min_samples_split：划分子节点需要的最少样本数据条数
    min_samples_leaf：叶子节点最少样本树
"""
model = DecisionTreeClassifier(criterion="gini")    # 创建模型实例对象
model.fit(x_train, y_train)                         # 模型训练

# 7.模型评估
# 7.1 使用测试集进行模型预测
y_pre = model.predict(x_test)
print(f"预测准确率为：{accuracy_score(y_test, y_pre)}")
print(f"精确率为：{precision_score(y_test, y_pre)}")
print(f"召回率为：{recall_score(y_test, y_pre)}")
print(f"F1分数为：{f1_score(y_test, y_pre)}")
print(f"预测报告：{classification_report(y_test, y_pre)}")

# 8.决策树图形绘制
plt.figure(figsize=(20, 20), dpi=200)
"""
    decision_tree：决策树算法模型实例对象
    max_depth：设置决策树显示的最大深度。对算法模型没影响
    feature_names：特征字段名称
    class_names：分类类别名称
    filled：决策树的节点是否进行填充展示
"""
plot_tree(
    decision_tree=model,
    max_depth=5,
    feature_names=x.columns.tolist(),
    class_names=["no_Survived", "yes_Survived"],
    filled=True
    )
plt.savefig("tree.jpg")
plt.show()